#pragma once

#include <thrust/scan.h>

template <typename Container1,
          typename Container2>
struct InclusiveScan
{
  Container1 A;
  Container2 B;

  template <typename Range1, typename Range2>
  InclusiveScan(const Range1& X, const Range2& Y)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end())
  {}

  void operator()(void)
  {
    thrust::inclusive_scan(A.begin(), A.end(), B.begin());
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename BinaryOperation>
struct InclusiveScanBinaryOp
{
  Container1 A;
  Container2 B;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2>
  InclusiveScanBinaryOp(const Range1& X, const Range2& Y, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::inclusive_scan(A.begin(), A.end(), B.begin(), binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2>
struct ExclusiveScan
{
  Container1 A;
  Container2 B;

  template <typename Range1, typename Range2>
  ExclusiveScan(const Range1& X, const Range2& Y)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end())
  {}

  void operator()(void)
  {
    thrust::exclusive_scan(A.begin(), A.end(), B.begin());
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename T>
struct ExclusiveScanInit
{
  Container1 A;
  Container2 B;
  T init;

  template <typename Range1, typename Range2>
  ExclusiveScanInit(const Range1& X, const Range2& Y, T init)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      init(init)
  {}

  void operator()(void)
  {
    thrust::exclusive_scan(A.begin(), A.end(), B.begin(), init);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename T,
          typename BinaryOperation>
struct ExclusiveScanInitBinaryOp
{
  Container1 A;
  Container2 B;
  T init;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2>
  ExclusiveScanInitBinaryOp(const Range1& X, const Range2& Y, T init, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      init(init),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::exclusive_scan(A.begin(), A.end(), B.begin(), init, binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3>
struct InclusiveScanByKey
{
  Container1 A;
  Container2 B;
  Container3 C;

  template <typename Range1, typename Range2, typename Range3>
  InclusiveScanByKey(const Range1& X, const Range2& Y, const Range3& Z)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end())
  {}

  void operator()(void)
  {
    thrust::inclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin());
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename BinaryPredicate>
struct InclusiveScanByKeyPred
{
  Container1 A;
  Container2 B;
  Container3 C;
  BinaryPredicate pred;

  template <typename Range1, typename Range2, typename Range3>
  InclusiveScanByKeyPred(const Range1& X, const Range2& Y, const Range3& Z, BinaryPredicate pred)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      pred(pred)
  {}

  void operator()(void)
  {
    thrust::inclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin(), pred);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};


template <typename Container1,
          typename Container2,
          typename Container3,
          typename BinaryPredicate,
          typename BinaryOperation>
struct InclusiveScanByKeyPredBinaryOp
{
  Container1 A;
  Container2 B;
  Container3 C;
  BinaryPredicate pred;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2, typename Range3>
  InclusiveScanByKeyPredBinaryOp(const Range1& X, const Range2& Y, const Range3& Z, BinaryPredicate pred, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      pred(pred),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::inclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin(), pred, binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3>
struct ExclusiveScanByKey
{
  Container1 A;
  Container2 B;
  Container3 C;

  template <typename Range1, typename Range2, typename Range3>
  ExclusiveScanByKey(const Range1& X, const Range2& Y, const Range3& Z)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end())
  {}

  void operator()(void)
  {
    thrust::exclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin());
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename T>
struct ExclusiveScanByKeyInit
{
  Container1 A;
  Container2 B;
  Container3 C;
  T init;

  template <typename Range1, typename Range2, typename Range3>
  ExclusiveScanByKeyInit(const Range1& X, const Range2& Y, const Range3& Z, T init)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      init(init)
  {}

  void operator()(void)
  {
    thrust::exclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin(), init);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename T,
          typename BinaryPredicate>
struct ExclusiveScanByKeyInitPred
{
  Container1 A;
  Container2 B;
  Container3 C;
  T init;
  BinaryPredicate pred;

  template <typename Range1, typename Range2, typename Range3>
  ExclusiveScanByKeyInitPred(const Range1& X, const Range2& Y, const Range3& Z, T init, BinaryPredicate pred)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      init(init),
      pred(pred)
  {}

  void operator()(void)
  {
    thrust::exclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin(), init, pred);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename T,
          typename BinaryPredicate,
          typename BinaryOperation>
struct ExclusiveScanByKeyInitPredBinaryOp
{
  Container1 A;
  Container2 B;
  Container3 C;
  T init;
  BinaryPredicate pred;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2, typename Range3>
  ExclusiveScanByKeyInitPredBinaryOp(const Range1& X, const Range2& Y, const Range3& Z, T init, BinaryPredicate pred, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      init(init),
      pred(pred),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::exclusive_scan_by_key(A.begin(), A.end(), B.begin(), C.begin(), init, pred, binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

